from typing import Dict, Sequence, Any

from mcp.types import TextContent, Tool

from core.exceptions import SQLExecutionError
from tools.base import ToolsBase
from utils.execute_sql_util import ExecuteSqlUtil


class ExecuteSQLTool(ToolsBase):
    """SQL 执行工具类
    
    该工具专门用于执行由 sql_creator 生成的 SQL 语句。它不应该直接执行模型自行构造的 SQL，
    而应在 sql_creator 成功生成 SQL 后调用此工具进行执行。支持执行多条 SQL 语句（以分号分隔）。
    """

    # 工具名称
    name = "execute_sql"
    # 工具描述，包含中英文说明和使用规范
    description = ("该工具仅用于执行由 sql_creator 生成的 SQL。禁止在此工具中直接输入模型自行构造或推测的 SQL 语句。"
                   "通常在 sql_creator 成功生成 SQL 后，才应调用此工具进行执行。"
                    "SQL 执行工具，用于在数据库上执行已生成的 SQL 语句（支持多条，以分号分隔）。"
                   "This tool is only for executing SQL generated by sql_creator. It is prohibited to input SQL statements "
                   "constructed or inferred by the model directly into this tool. This tool should typically be called for "
                   "execution only after sql_creator successfully generates SQL. SQL execution tool, used to execute generated SQL "
                   "statements on the database (supports multiple statements separated by semicolons)."
                   )


    def get_tool_description(self) -> Tool:
        """获取工具的详细描述信息
        
        返回一个Tool对象，包含工具名称、描述和输入参数模式。
        该描述用于MCP服务识别和调用工具。
        
        Returns:
            Tool: 包含工具描述信息的对象
        """
        return Tool(
            name=self.name,
            description=self.description,
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "要执行的SQL语句"
                    },
                    "pool_name": {
                        "type": "string",
                        "description": "线程池名称,若没有指定默认是default"
                    }
                },
                "required": ["query"]
            }
        )


    async def run_tool(self, arguments: Dict[str, Any]) -> Sequence[TextContent]:
        """执行SQL工具
        
        执行传入的SQL语句，并返回执行结果。
        
        Args:
            arguments: 包含执行参数的字典
                - query (str): 要执行的SQL语句
                - pool_name (str, optional): 数据库连接池名称，默认为"default"

        Returns:
            Sequence[TextContent]: 包含执行结果的文本内容序列
            
        Raises:
            SQLExecutionError: 当SQL执行出现错误时抛出
        """
        # 检查参数中是否包含必需的query字段
        if "query" not in arguments:
            return [TextContent(type="text", text="错误: 缺少查询语句")]

        # 获取要执行的SQL语句
        query = arguments["query"]

        # 获取连接池名称，如果没有指定则使用默认值"default"
        if "pool_name" not in arguments:
            pool_name = "default"
        else:
            pool_name = arguments["pool_name"]

        try:
            # 执行多条SQL语句
            sql_results = ExecuteSqlUtil.execute_multiple_statements(pool_name, query)

            # 格式化执行结果
            results = []
            for result in sql_results:
                formatted_result = ExecuteSqlUtil.format_result(result)
                results.append(formatted_result)

            # 将所有结果用分隔符连接并返回
            return [TextContent(type="text", text="\n---\n".join(results))]

        # 捕获SQL执行错误并返回错误信息
        except SQLExecutionError as e:
            return [TextContent(type="text", text=str(e))]